package com.syf.starter.core.netty.tool.spi.impl;

import com.google.common.util.concurrent.RateLimiter;
import com.syf.starter.core.netty.tool.spi.SyfServerInvokerShape;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;

/**
 * 公平的令牌桶限流器
 * */
@Slf4j
public class FairTokenBucketShape implements SyfServerInvokerShape {
    private RateLimiter limiter;

    private BlockingQueue<traffic> requestQueue = new LinkedBlockingQueue<>(1024);

    private Thread shapeThread;

    private boolean isStop = false;

    public FairTokenBucketShape(Integer threadNum, Float avgResponseSeconds) {
        this.limiter = RateLimiter.create(setRate(threadNum, avgResponseSeconds).doubleValue());

        this.shapeThread = new Thread(() -> {
            try {
                while (!isStop) {
                    traffic traffic = requestQueue.take();

                    // 排队等待期间，就超时了
                    if (traffic.getTimeout() == 0L) {
                        continue;
                    }
                    
                    // 能否获得令牌
                    if (limiter.tryAcquire(traffic.getTimeout(), TimeUnit.MILLISECONDS)) {
                        traffic.getLatch().countDown();
                    }
                }
            } catch (InterruptedException e) {
                log.info("公平令牌桶-流量整形器被关闭");
                isStop = true;
            }
        });

        shapeThread.start();
    }

    @Override
    public boolean tryAcquire(long timeout, TimeUnit unit) {
        CountDownLatch latch = new CountDownLatch(1);
        traffic traffic = new traffic(latch, timeout, unit);
        try {
            requestQueue.put(traffic);

            if (latch.await(timeout, unit)) {
                return true;
            }
            return false;
        } catch (InterruptedException e) {
            log.warn("FairTokenBucketShape中队列已满，service actuator is too busy！");
            return false;
        }
    }

    @Override
    public void acquire() {
        limiter.acquire();
    }

    @Override
    public void shutdown() {
        isStop = true;
        shapeThread.interrupt();
    }

    public static class traffic {
        private CountDownLatch latch;

        private Long timeout;

        private Long timeStamp;

        public traffic(CountDownLatch latch, Long timeout, TimeUnit unit) {
            this.latch = latch;
            this.timeout = unit.toMillis(timeout);
            timeStamp = System.currentTimeMillis();
        }

        public CountDownLatch getLatch() {
            return latch;
        }

        public Long getTimeout() {
            return Math.max(0, timeout - System.currentTimeMillis() + timeStamp);
        }

        public Long getTimeStamp() {
            return timeStamp;
        }
    }
}
